import os
import os
import logging
from contextlib import asynccontextmanager
from typing import Dict, Any, Optional, List
from dataclasses import dataclass
from functools import lru_cache

from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from openai import AsyncOpenAI
import asyncio

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 数据模型
class TestRequest(BaseModel):
    class_path: str = Field(..., description="完整的Java类路径，如：com.example.service.UserService")
    method_code: str = Field(..., description="需要生成测试的Java方法代码")

class TestResponse(BaseModel):
    test_code: str = Field(..., description="生成的JUnit测试代码")
    provider: str = Field(..., description="使用的AI提供商")
    model: str = Field(..., description="使用的AI模型")

class ProviderSwitchRequest(BaseModel):
    provider: str = Field(..., description="要切换到的AI提供商")

class ConfigResponse(BaseModel):
    current_provider: str
    current_model: str
    available_providers: List[str]

class HealthResponse(BaseModel):
    status: str
    provider: str
    model: str

@dataclass
class ModelConfig:
    api_key: Optional[str]
    base_url: str
    model: str
    provider_name: str
    requires_auth: bool = True

# 模型配置管理
class ModelManager:
    def __init__(self):
        self.configs = self._load_configs()
        self.current_provider = os.getenv("AI_PROVIDER", "deepseek")
        self.clients = {}
        self._initialize_clients()
    
    def _load_configs(self) -> Dict[str, ModelConfig]:
        """加载所有模型配置"""
        configs = {}
        
        # DeepSeek
        configs["deepseek"] = ModelConfig(
            api_key=os.getenv("DEEPSEEK_API_KEY"),
            base_url=os.getenv("DEEPSEEK_API_BASE", "https://api.deepseek.com/v1"),
            model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
            provider_name="DeepSeek"
        )
        
        # OpenAI
        configs["openai"] = ModelConfig(
            api_key=os.getenv("OPENAI_API_KEY"),
            base_url=os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"),
            model=os.getenv("OPENAI_MODEL", "gpt-4o-mini"),
            provider_name="OpenAI"
        )
        
        # Azure OpenAI
        configs["azure"] = ModelConfig(
            api_key=os.getenv("AZURE_API_KEY"),
            base_url=os.getenv("AZURE_API_BASE", ""),
            model=os.getenv("AZURE_MODEL", "gpt-4o-mini"),
            provider_name="Azure OpenAI"
        )
        
        # Claude
        configs["claude"] = ModelConfig(
            api_key=os.getenv("CLAUDE_API_KEY"),
            base_url=os.getenv("CLAUDE_API_BASE", "https://api.anthropic.com"),
            model=os.getenv("CLAUDE_MODEL", "claude-3-5-sonnet-20241022"),
            provider_name="Claude"
        )
        
        # Gemini
        configs["gemini"] = ModelConfig(
            api_key=os.getenv("GEMINI_API_KEY"),
            base_url=os.getenv("GEMINI_API_BASE", "https://generativelanguage.googleapis.com/v1beta"),
            model=os.getenv("GEMINI_MODEL", "gemini-1.5-flash"),
            provider_name="Gemini"
        )
        
        # 清华智谱
        configs["zhipu"] = ModelConfig(
            api_key=os.getenv("ZHIPU_API_KEY"),
            base_url=os.getenv("ZHIPU_API_BASE", "https://open.bigmodel.cn/api/paas/v4"),
            model=os.getenv("ZHIPU_MODEL", "glm-4"),
            provider_name="清华智谱"
        )
        
        # OpenRouter（支持免费模型）
        openrouter_key = os.getenv("OPENROUTER_API_KEY")
        configs["openrouter"] = ModelConfig(
            api_key=openrouter_key if openrouter_key else "sk-optional",  # 可选
            base_url=os.getenv("OPENROUTER_API_BASE", "https://openrouter.ai/api/v1"),
            model=os.getenv("OPENROUTER_MODEL", "microsoft/deepseek-r1:free"),
            provider_name="OpenRouter",
            requires_auth=bool(openrouter_key)  # 如果没有key，使用免费模型
        )
        
        return configs
    
    def _initialize_clients(self):
        """初始化配置的模型客户端"""
        # 只初始化当前配置的提供商模型
        if self.current_provider in self.configs:
            config = self.configs[self.current_provider]
            if self._is_config_valid(config):
                try:
                    client_kwargs = {
                        'base_url': config.base_url,
                        'timeout': 60.0
                    }
                    
                    # 对于不需要认证的开放模型，可以留空或填占位符
                    if config.requires_auth and config.api_key:
                        client_kwargs['api_key'] = config.api_key
                    elif not config.requires_auth:
                        client_kwargs['api_key'] = "free-model"
                    else:
                        return
                    
                    self.clients[self.current_provider] = AsyncOpenAI(**client_kwargs)
                    logger.info(f"已初始化 {config.provider_name} 客户端 - 模型: {config.model}")
                except Exception as e:
                    logger.warning(f"初始化 {config.provider_name} 失败: {e}")
    
    def _is_config_valid(self, config: ModelConfig) -> bool:
        """检查配置是否有效"""
        if config.provider_name == "Azure OpenAI" and not config.api_key:
            return False
        if config.requires_auth and not config.api_key:
            return False
        if not config.base_url:
            return False
        return True
    
    def get_available_providers(self) -> List[str]:
        """获取可用的提供商列表"""
        return list(self.clients.keys())
    
    def switch_provider(self, provider: str) -> bool:
        """切换AI提供商"""
        if provider in self.clients:
            self.current_provider = provider
            return True
        return False
    
    def get_current_config(self) -> ModelConfig:
        """获取当前提供商配置"""
        return self.configs.get(self.current_provider)
    
    def get_client(self) -> AsyncOpenAI:
        """获取当前客户端"""
        return self.clients.get(self.current_provider)

# 全局模型管理器
model_manager = ModelManager()

# 生命周期管理
@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("启动Java单元测试生成器...")
    logger.info(f"当前AI提供商: {model_manager.current_provider}")
    yield
    logger.info("正在关闭应用...")

# FastAPI应用
app = FastAPI(
    title="Java单元测试生成器",
    description="基于大语言模型的Java单元测试代码生成工具，支持多种AI提供商",
    version="2.1.0",
    lifespan=lifespan
)

# CORS配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 静态文件
app.mount("/static", StaticFiles(directory="app/static"), name="static")

# 路由
@app.get("/", response_class=HTMLResponse)
async def read_root():
    return FileResponse("app/templates/index.html")

@app.post("/api/generate-test", response_model=TestResponse)
async def generate_test(request: TestRequest):
    """生成Java单元测试代码"""
    try:
        client = model_manager.get_client()
        config = model_manager.get_current_config()
        
        if not client or not config:
            raise HTTPException(
                status_code=500,
                detail=f"当前AI提供商 {model_manager.current_provider} 不可用"
            )
        
        # 生成提示词
        prompt = f"""你是一个专业的Java测试工程师，请为以下Java方法生成完整的JUnit 4测试代码。

类路径：{request.class_path}

方法代码：
{request.method_code}

要求：
1. 使用JUnit 4和Mockito框架
2. 包含正常流程和边界条件的测试
3. 添加适当的Mock对象设置
4. 使用断言验证预期结果
5. 包含异常情况的测试
6. 添加清晰的测试注释
7. 遵循Java命名规范

请只返回完整的Java测试代码，不要添加解释。"""

        # 调用AI生成测试代码
        response = await client.chat.completions.create(
            model=config.model,
            messages=[
                {"role": "system", "content": "你是一个专业的Java测试工程师，专门生成高质量的JUnit测试代码。"},
                {"role": "user", "content": prompt}
            ],
            temperature=0.1,
            max_tokens=4000
        )
        
        test_code = response.choices[0].message.content.strip()
        
        # 清理代码格式
        if test_code.startswith("```java"):
            test_code = test_code[7:]
        if test_code.startswith("```"):
            test_code = test_code[3:]
        if test_code.endswith("```"):
            test_code = test_code[:-3]
        
        return TestResponse(
            test_code=test_code.strip(),
            provider=model_manager.current_provider,
            model=config.model
        )
        
    except Exception as e:
        logger.error(f"生成测试代码失败: {e}")
        raise HTTPException(
            status_code=500,
            detail=f"生成测试代码失败: {str(e)}"
        )

@app.get("/api/config", response_model=ConfigResponse)
async def get_config():
    """获取当前配置信息"""
    config = model_manager.get_current_config()
    return ConfigResponse(
        current_provider=model_manager.current_provider,
        current_model=config.model if config else "",
        available_providers=model_manager.get_available_providers()
    )

@app.post("/api/switch-provider")
async def switch_provider(request: ProviderSwitchRequest):
    """切换AI提供商"""
    if request.provider not in model_manager.get_available_providers():
        raise HTTPException(
            status_code=400,
            detail=f"不支持的AI提供商: {request.provider}. 可用提供商: {model_manager.get_available_providers()}"
        )
    
    if model_manager.switch_provider(request.provider):
        config = model_manager.get_current_config()
        return {
            "success": True,
            "provider": request.provider,
            "model": config.model if config else ""
        }
    else:
        raise HTTPException(
            status_code=500,
            detail="切换AI提供商失败"
        )

@app.get("/api/health", response_model=HealthResponse)
async def health_check():
    """健康检查"""
    config = model_manager.get_current_config()
    return HealthResponse(
        status="healthy",
        provider=model_manager.current_provider,
        model=config.model if config else "unknown"
    )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "app.main:app",
        host=os.getenv("HOST", "127.0.0.1"),
        port=int(os.getenv("PORT", 8000)),
        reload=os.getenv("DEBUG", "false").lower() == "true"
    )